#include "compensator.hpp"

#include "opencv2/opencv.hpp"
#include "spdlog/spdlog.h"

namespace {

const double kG = 9.80665;
const cv::Scalar kGREEN(0., 255., 0.);
const cv::Scalar kRED(0., 0., 255.);
const cv::Scalar kYELLOW(0., 255., 255.);
const auto kCV_FONT = cv::FONT_HERSHEY_SIMPLEX;

}  // namespace

/**
 * @brief Angle θ required to hit coordinate (x, y)
 *
 * {\displaystyle \tan \theta ={\left({\frac {v^{2}\pm {\sqrt
 * {v^{4}-g(gx^{2}+2yv^{2})}}}{gx}}\right)}}
 *
 * @param target 目标坐标
 * @return double 出射角度
 */
double Compensator::SolveSurfaceLanchAngle(cv::Point2f target) {
  const double v_2 = pow(ballet_speed_, 2);
  const double up_base =
      std::sqrt(std::pow(ballet_speed_, 4) -
                kG * (kG * std::pow(target.x, 2) + 2 * target.y * v_2));
  const double low = kG * target.x;
  const double ans1 = std::atan((v_2 + up_base) / low);
  const double ans2 = std::atan((v_2 - up_base) / low);

  if (std::isnan(ans1)) return std::isnan(ans2) ? 0. : ans2;
  if (std::isnan(ans2)) return std::isnan(ans1) ? 0. : ans1;
  return std::min(ans1, ans2);
}

double Compensator::SolveOffset(float distance, float pitch) {
  const float gun_cam_distance = 0;
  const float cam_target_height = distance * sin(pitch);
  const float gun_target_height = cam_target_height + gun_cam_distance;  //高度y
  const float gun_target_dis = distance * cos(pitch);  //距离x

  return SolveSurfaceLanchAngle(cv::Point2f(gun_target_dis, gun_target_height));
}

void Compensator::VisualizePnp(Armor& armor, const cv::Mat& output,
                               bool add_lable) {
  std::vector<cv::Point2f> out_points;
  std::vector<cv::Point3f> world_points;
  const cv::Matx33f rmat = cv::Matx33f::eye();
  cv::Mat rvec;
  armor.SurfaceCenter();
  cv::Rodrigues(rmat, rvec);
  cv::Mat tvec = cv::Mat::zeros(3, 1, CV_32F);
  world_points.push_back(cv::Point3f(0, 0, 0));
  world_points.push_back(cv::Point3f(10, 0, 0));  // x
  world_points.push_back(cv::Point3f(0, 10, 0));  // y
  world_points.push_back(cv::Point3f(0, 0, 10));  // z

  //将world坐标中的原点以及坐标轴方向映射到2D
  cv::projectPoints(world_points, rvec, tvec, cv::Matx33f(cam_mat_),
                    distor_coff_, out_points);

  cv::line(output, out_points[0], out_points[1], kYELLOW, 2);  // x yellow
  cv::line(output, out_points[0], out_points[2], kGREEN, 4);   // y green
  cv::line(output, out_points[0], out_points[3], kRED, 6);     // z red

  if (add_lable) {
    cv::putText(output, "PNP", out_points[0], kCV_FONT, 1.0, kYELLOW);
    cv::putText(output, "x", out_points[1], kCV_FONT, 1.0, kGREEN);
    cv::putText(output, "y", out_points[2], kCV_FONT, 1.0, kGREEN);
    cv::putText(output, "z", out_points[3], kCV_FONT, 1.0, kGREEN);
  }
}

Compensator::Compensator() { SPDLOG_TRACE("Constructed."); }

Compensator::Compensator(const std::string& cam_mat_path) {
  SPDLOG_TRACE("Constructed.");
  LoadCameraMat(cam_mat_path);
}

Compensator::~Compensator() { SPDLOG_TRACE("Destructed."); }

void Compensator::LoadCameraMat(const std::string& path) {
  cv::FileStorage fs(path,
                     cv::FileStorage::READ | cv::FileStorage::FORMAT_JSON);

  if (fs.isOpened()) {
    cam_mat_ = fs["cam_mat"].mat();
    distor_coff_ = fs["distor_coff"].mat();
    if (cam_mat_.empty() && distor_coff_.empty()) {
      SPDLOG_ERROR("Can not load cali data.");
    } else {
      SPDLOG_DEBUG("Loaded cali data.");
    }
  } else {
    SPDLOG_ERROR("Can not open file: '{}'", path);
  }
}

cv::Point3f Compensator::EstimateWorldCoord(Armor& armor, double& distance) {
  cv::Mat rot_vec, trans_vec;
  const double gun_cam_distance = 0.;  //枪口和摄像头距离y
  cv::solvePnP(armor.ModelVertices(), armor.SurfaceVertices(), cam_mat_,
               distor_coff_, rot_vec, trans_vec, false, cv::SOLVEPNP_ITERATIVE);
  // trans_vec = trans_vec * (5. / 3.);
  armor.SetRotVec(rot_vec), armor.SetTransVec(trans_vec);
  trans_vec.at<double>(1, 0) -= gun_cam_distance;
  distance = cv::norm(trans_vec);

  return cv::Vec3f(trans_vec);
}

void Compensator::Apply(std::vector<Armor>& armors, const cv::Mat& frame,
                        float ballet_speed, float robot_pitch) {
  const cv::Point2f frame_center(frame.cols / 2, frame.rows / 2);
  ballet_speed_ = ballet_speed * 1000.;  //同步射速,单位mm
  for (auto& armor : armors) {
    double distance = 0;
    common::Euler aiming_eulr;
    auto trans_vec = EstimateWorldCoord(armor, distance);  //相对于相机
    SPDLOG_WARN("distance is {}", distance);
    SPDLOG_WARN("vec is {}, {}, {}", trans_vec.x, trans_vec.y, trans_vec.z);

    float pitch = 0.f;
    if (distance > 5000) {  //距离太远
      const double fx = cam_mat_.at<double>(0, 0);
      const double fy = cam_mat_.at<double>(1, 1);
      const double cx = cam_mat_.at<double>(0, 2);
      const double cy = cam_mat_.at<double>(1, 2);

      cv::Point2f pnt;
      std::vector<cv::Point2f> in, out;
      in.push_back(armor.SurfaceCenter());

      //对像素点去畸变处理
      cv::undistortPoints(in, out, cam_mat_, distor_coff_, cv::noArray(),
                          cam_mat_);
      pnt = out.front();
      //去畸变后的比值
      const double rxNew = (pnt.x - cx) / fx;
      const double ryNew = (pnt.y - cy) / fy;
      aiming_eulr.yaw = atan(rxNew);
      pitch = 0.5 * atan(ryNew);
      aiming_eulr.yaw = (frame_center.x - armor.SurfaceCenter().x) /
                        frame.cols * 47.9 / 180 * CV_PI * 0.2;
    } else {  //一般距离
      const double x_pos = trans_vec.x, y_pos = trans_vec.y,
                   z_pos = trans_vec.z;

      SPDLOG_WARN("x,y,z is {},{},{}", x_pos, y_pos, z_pos);
      const double tan_pitch = y_pos / sqrt(x_pos * x_pos + z_pos * z_pos);
      pitch = 0.7 * atan(tan_pitch);
      aiming_eulr.yaw = atan(x_pos / z_pos);
      aiming_eulr.yaw = (frame_center.x - armor.SurfaceCenter().x) /
                        frame.cols * 47.9 / 180 * CV_PI * 0.5;
    }

    // aiming_eulr.pitch = pitch;
    aiming_eulr.pitch = SolveOffset(distance, pitch + robot_pitch);
    aiming_eulr.roll = 0.f;
    SPDLOG_WARN("pitch is {}", pitch);
    SPDLOG_WARN("armor center {}, {}", armor.SurfaceCenter().x,
                armor.SurfaceCenter().y);  // y>240
    SPDLOG_WARN("armor.center > frame.center {}",
                armor.SurfaceCenter().y > frame_center.y);
    SPDLOG_WARN("aiming_eulr is yaw {}, pitch {}", aiming_eulr.yaw,
                aiming_eulr.pitch);
    armor.SetAimEuler(aiming_eulr);
  }
  std::sort(armors.begin(), armors.end(),
            [frame_center](Armor& armor1, Armor& armor2) {
              return cv::norm(armor1.SurfaceCenter() - frame_center) <
                     cv::norm(armor2.SurfaceCenter() - frame_center);
            });
}

void Compensator::VisualizeResult(std::vector<Armor>& armors,
                                  const cv::Mat& output, int verbose) {
  for (auto& armor : armors) {
    VisualizePnp(armor, output, verbose > 1);
    // VisualizeSolveAngle(armor, output, verbose > 2);
  }
}
